package personal.wang;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import org.apache.commons.io.FileUtils;
import org.apache.kafka.clients.admin.*;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartitionInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ExecutionException;

import static com.alibaba.fastjson2.JSONWriter.Feature.PrettyFormat;

/**
 * @className: AdminClientTest
 * @Description:
 * @Author: wangyifei
 * @Date: 2025/4/22 9:53
 */
public class RecopyReplicas {
    private static Logger logger = LoggerFactory.getLogger(RecopyReplicas.class);
    private static String KEY_PARTITIONS = "partitions";
    private static Integer index = 0;

    public static void run(String[] args) {
        Properties props = new Properties();
        props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "kafka8:9092,kafka9:9092,kafka10:9092,kafka11:9092,kafka12:9092");
        props.put("request.timeout.ms", 600000);
        String filesPath = "/home/bigdata/kafka_about_shell/reassign/";
        int[] leftBrokerId = {1,2,3,7,8,9,10,11,12};
        int[] deleteBrokerId = {4,5,6,13,14,15,16,17};

        try (AdminClient client = AdminClient.create(props)) {
            ListTopicsResult listTopicsResult = client.listTopics();
            Set<String> topics = listTopicsResult.names().get();
            for (String topic : topics) {
                Map<String, Object> reassignedReplica = reassignReplica();
                JSONArray partitions = (JSONArray) reassignedReplica.get(KEY_PARTITIONS);
                DescribeTopicsResult describeTopicsResult = client.describeTopics(Collections.singletonList(topic));
                Map<String, TopicDescription> descriptionMap = describeTopicsResult.all().get();
                // 找出在 deleteBrokerId 的副本信息
                descriptionMap.forEach((tpc , topicDesc) -> {
                    for (TopicPartitionInfo partition : topicDesc.partitions()) {
                        int[] allArray = partition.replicas().stream().mapToInt(Node::id).sorted().toArray();
                        int[] toDeleteReplica = Arrays.stream(allArray).filter(x -> Arrays.binarySearch(deleteBrokerId, x) >= 0).sorted().toArray();
                        if(toDeleteReplica.length > 0){
                            Map<String,Object> p = new HashMap<>();
                            p.put("topic" , tpc);
                            p.put("partition" , partition.partition());
                            p.put("replicas" , Arrays.stream(allArray).map(x -> {
                                if(Arrays.binarySearch(toDeleteReplica , x) >= 0){
                                    return replaceReplica(leftBrokerId , allArray);
                                }
                                return x ;
                            }).toArray());
                            partitions.add(p);
                        }
                    }
                });
                if(!partitions.isEmpty()){
                    String jsonString = JSONObject.toJSONString(reassignedReplica);
                    FileUtils.writeStringToFile(new File(filesPath + topic + ".json") , jsonString , StandardCharsets.UTF_8.displayName());
                }
            }
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private static int replaceReplica(int[] leftBrokerId , int[] all) {
        index = ( index + 1 )%leftBrokerId.length;
        while(Arrays.binarySearch(all , leftBrokerId[index]) >= 0){
            index = ( index + 1 )%leftBrokerId.length;
        }
        return leftBrokerId[index];
    }

    public static Map<String,Object> reassignReplica(){
        Map<String , Object> map = new HashMap<>();
        map.put("version",1);
        map.put("partitions",new JSONArray());
        return map;
    }

    public static void main(String[] args) {
        run(args);
    }

    public static void exe(String[] args) {
        Map<String, Object> stringObjectMap = reassignReplica();
        JSONArray o = (JSONArray)stringObjectMap.computeIfAbsent("partitions", x -> new JSONArray());
        JSONObject obj = new JSONObject();
        obj.put("topic","wo-video-logback");
        obj.put("partition",1);
        obj.put("replicas",new int[]{1,2,3});
        o.add(obj);
        System.out.println(JSONObject.toJSONString(stringObjectMap, PrettyFormat));

        System.out.println(Arrays.binarySearch(new int[]{1, 9, 2, 3, 2}, 2));
        System.out.println(Arrays.binarySearch(new int[]{1, 9, 2, 3, 2}, 2));
    }

}
